import torch
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm
from bert_config_movie import MAX_LEN, device, model_dir
from fine_tuning_data_formatting import tokenizer
import pandas as pd
from fine_tuning import BertClass
import numpy as np

import warnings

warnings.filterwarnings("ignore")


# 定义一个自定义数据集类QueryData，继承自Dataset类
class QueryData(Dataset):
    def __init__(self, dataframe, tokenizer, max_len):
        # 初始化数据集的BertTokenizer
        self.tokenizer = tokenizer
        self.data = dataframe
        # 从DataFrame中获取查询文本列
        self.text = dataframe['query']
        #         self.targets = self.data.citation_influence_label
        # 初始化数据集的最大文本长度
        self.max_len = max_len

    # 定义数据集的长度，即数据集中样本的数量
    def __len__(self):
        return len(self.text)

    # 定义数据集的获取方法，根据索引获取具体样本
    def __getitem__(self, index):
        text = str(self.text[index])
        # 对查询文本进行分词，去除多余的空格
        text = " ".join(text.split())

        # 使用self.tokenizer.encode_plus()方法对文本进行编码，并生成模型的输入张量
        # 包含以下信息
        # input_ids：编码后的文本，表示为数字序列。
        # attention_mask：注意力掩码，用于指示哪些令牌在模型的注意力机制中被考虑，哪些令牌被掩盖。
        # token_type_ids：用于区分输入中不同句子或段落的令牌类型。在这个任务中，token_type_ids全为0，因为只有一个文本序列。
        inputs = self.tokenizer.encode_plus(
            text,  # 要编码的文本
            None,  # 在此情况下，表示不存在关联的另一个文本。在这个任务中，文本之间没有关联，因此使用None
            add_special_tokens=True,  # 表示在编码时添加特殊标记，例如[CLS]和[SEP]。
            max_length=self.max_len,  # 表示编码后的序列的最大长度。
            pad_to_max_length=True,  # 如果编码后的序列长度小于max_length，则在序列末尾填充0，使其长度与max_length相同。
            return_token_type_ids=True  # 用于区分输入中的不同句子或段落。在这个任务中，由于输入没有关联的另一个文本，token_type_ids全为0。
        )
        ids = inputs['input_ids']
        mask = inputs['attention_mask']
        token_type_ids = inputs["token_type_ids"]
        # 返回字典
        return {
            'ids': torch.tensor(ids, dtype=torch.long),
            'mask': torch.tensor(mask, dtype=torch.long),
            'token_type_ids': torch.tensor(token_type_ids, dtype=torch.long)
            #             'targets': torch.tensor(self.targets[index], dtype=torch.float)
        }


# 测试函数，用于在测试集上进行模型预测
def test(model, testing_loader):
    res = []
    # 将BERT模型设置为评估模式
    model.eval()
    # 使用torch.no_grad()上下文管理器，禁用梯度计算
    with torch.no_grad():
        for _, data in tqdm(enumerate(testing_loader, 0)):
            # 将输入数据的ids张量移动到指定的计算设备
            ids = data['ids'].to(device, dtype=torch.long)
            mask = data['mask'].to(device, dtype=torch.long)
            token_type_ids = data['token_type_ids'].to(device, dtype=torch.long)
            #             targets = data['targets'].to(device, dtype = torch.long)
            # 通过BERT模型进行前向传播，得到预测输出。
            outputs = model(ids, mask, token_type_ids)
            my_outputs = outputs.tolist()
            score = np.exp(my_outputs[0]) / np.sum(np.exp(my_outputs[0]), axis=0)
            idx = np.argmax(score)
            # 获取预测输出中每个样本预测得分最高的类别索引。
            big_val, big_idx = torch.max(outputs, dim=1)
            # 将预测结果添加到列表res中。
            res.extend(big_idx.tolist())
            res.append(round(score[idx] * 100, 2))
            print(res)
    return res


if __name__ == '__main__':
    # 加载训练好的模型
    model = torch.load(model_dir, map_location="cpu")
    # model = BertClass()
    # model.load_state_dict(torch.load(model_dir, map_location="cpu"), strict=False)
    while True:
        print("=" * 50)
        test_movie_review = input("Please enter your movie review > ")
        if test_movie_review == "#" or test_movie_review == "q":
            print("Exit the application !")
            exit(0)
        if test_movie_review == "":
            print("Please input something ~")
            continue
        test_data = pd.DataFrame([test_movie_review], columns=["query"])

        # 创建测试数据集
        data_to_test = QueryData(test_data[['query']], tokenizer, MAX_LEN)

        # 测试参数设置
        test_params = {'batch_size': 1,  # 批次大小
                       'shuffle': False,  # 是否打乱顺序
                       'num_workers': 0  # 运行的线程数量
                       }
        # 用于测试的数据加载器。
        testing_loader_f = DataLoader(data_to_test, **test_params)
        # 对测试集进行预测，得到预测结果
        res = test(model, testing_loader_f)
        if res[0] == 1:
            print("result > {} This is a positive movie review ~".format(res))
        else:
            print("result > {} This is a negative movie review !".format(res))
        print("=" * 50)
